// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::ffi::c_uint;

use smallvec::smallvec;
use tracing::info;

use cudax::nvml::*;
use xgpu_common::{
    api_name::ApiFuncName,
    ipc::message::{Argument, ArgumentFlag, Request},
};

use crate::{hook::nvml::NvmlApi, thread};

pub struct NvmlApiImpl;

impl NvmlApi for NvmlApiImpl {
    fn nvmlInit_v2(&self) -> nvmlReturn_t {
        info!("[Hooked] api_name: nvmlInit_v2");
        let req = Request::empty(ApiFuncName::FuncNvmlinitV2 as u64);
        thread::context()
            .agent()
            .invoke_api::<nvmlReturn_t>(req)
            .expect("call invoke_api failed")
    }

    fn nvmlShutdown(&self) -> nvmlReturn_t {
        info!("[Hooked] api_name: nvmlShutdown");
        let req = Request::empty(ApiFuncName::FuncNvmlshutdown as u64);
        thread::context()
            .agent()
            .invoke_api::<nvmlReturn_t>(req)
            .expect("call invoke_api failed")
    }

    fn nvmlDeviceGetCount_v2(&self, device_count: *mut c_uint) -> nvmlReturn_t {
        info!("[Hooked] api_name: nvmlDeviceGetCount_v2");
        let req = Request::with_args(
            ApiFuncName::FuncNvmldevicegetcountV2 as u64,
            smallvec![unsafe { Argument::from_mut_ptr(device_count, ArgumentFlag::ARG_OUT) }],
        );
        thread::context()
            .agent()
            .invoke_api::<nvmlReturn_t>(req)
            .expect("call invoke_api failed")
    }
}
